// package handlers 定义处理程序所在的包
package handlers

// 导入所需的包
import (
	"Lunar-Astral-API/local_server/config" // 导入项目配置包，用于获取配置信息
	"bytes"                                // 导入 bytes 包，用于处理字节切片
	"encoding/json"                        // 导入 json 包，用于 JSON 数据的编码和解码
	"fmt"                                  // 导入 fmt 包，用于格式化输入输出
	"io"                                   // 导入 io 包，提供基本的 I/O 操作接口
	"log"                                  // 导入 log 包，用于日志记录
	"net/http"                             // 导入 http 包，用于构建 HTTP 客户端和服务器
	"net/http/httputil"                    // 导入 httputil 包，提供 HTTP 协议的实用工具
	"net/url"                              // 导入 url 包，用于解析和处理 URL
	"strings"                              // 导入 strings 包，用于处理字符串
	"time"                                 // 导入 time 包，用于处理时间
)

// AgentRuleModelsHandler 处理获取模型列表的请求, 如果配置了 GGUF_Url，则将请求反向代理到该地址；否则返回本地模型列表。
// 该函数根据 HTTP 请求的方法和路径，尝试提取模型名称。
// 如果请求方法为 POST，则从请求体中提取 "model" 字段的值；
// 如果请求方法为 GET，则从 URL 路径中提取模型名称。
// 参数:
//   - w: http.ResponseWriter 类型，用于向客户端发送响应
//   - r: *http.Request 类型，代表 HTTP 请求
func AgentRuleModelsHandler(w http.ResponseWriter, r *http.Request) {
	// 检查是否配置了 GGUF_Url，如果已配置则将请求代理到该地址
	if config.ModelSystemPort != nil && *config.ModelSystemPort != 0 {
		// 构造目标 URL
		targetURL := fmt.Sprintf("http://localhost:%d", *config.ModelSystemPort)
		// 解析目标 URL
		target, err := url.Parse(targetURL)
		// 解析失败时返回 500 错误
		if err != nil {
			http.Error(w, "Failed to parse target URL", http.StatusInternalServerError)
			return
		}
		// 创建单主机反向代理
		proxy := httputil.NewSingleHostReverseProxy(target)
		// 将请求代理到目标地址
		proxy.ServeHTTP(w, r)
		// 代理完成后返回
		return
	}
	// 用于存储模型信息的切片
	models := []map[string]any{}
	// 加读锁，防止并发修改模型端口映射
	config.ModelMapMutex.RLock()
	// 函数结束时解锁
	defer config.ModelMapMutex.RUnlock()
	// 存储模型名称的切片
	var modelNames []string
	// 遍历模型端口映射，获取所有模型名称
	for modelName := range config.ModelPortMap {
		modelNames = append(modelNames, modelName)
	}
	// 遍历模型名称，构造模型信息
	for _, modelName := range modelNames {
		models = append(models, map[string]any{
			"id":       modelName,
			"object":   "model",
			"owned_by": "organization_owner",
		})
	}
	// 构造响应数据
	response := map[string]any{
		"object": "list",
		"data":   models,
	}
	// 设置响应头，指定返回数据为 JSON 格式
	w.Header().Set("Content-Type", "application/json")
	// 将响应数据编码为 JSON 并写入响应
	json.NewEncoder(w).Encode(response)
}

// AgentRuleHandler 处理与模型规则相关的请求, 如果配置了 GGUF_Url，则将请求反向代理到该地址；否则根据请求获取模型名称并代理到对应的端口。
// 该函数根据 HTTP 请求的方法和路径，尝试提取模型名称。
// 如果请求方法为 POST，则从请求体中提取 "model" 字段的值；
// 如果请求方法为 GET，则从 URL 路径中提取模型名称。
// 参数:
//   - w: http.ResponseWriter 类型，用于向客户端发送响应
//   - r: *http.Request 类型，代表 HTTP 请求
func AgentRuleHandler(w http.ResponseWriter, r *http.Request) {
	// 检查系统是否繁忙，如果已就绪的模型数量小于最大模型数量，返回系统繁忙响应
	if config.ModelReady < config.MaxModelAmount {
		serveBusyResponse(w)
		return
	}
	// 检查是否配置了模型系统端口，如果已配置则将请求代理到该端口
	if config.ModelSystemPort != nil && *config.ModelSystemPort != 0 {
		proxyToPort(w, r, *config.ModelSystemPort)
		return
	}
	// 从请求中提取模型名称
	modelName := extractModelName(r)
	// 如果未能提取到模型名称，返回 400 错误
	if modelName == "" {
		http.Error(w, "GGUF服务[ERROR] -> 无法从请求中提取模型名称", http.StatusBadRequest)
		return
	}
	// 根据模型名称获取对应的端口号
	port, exists := getModelPort(modelName)
	// 如果未找到对应的模型，返回 404 错误
	if !exists {
		http.Error(w, "GGUF服务[ERROR] -> 无法找到模型: "+modelName, http.StatusNotFound)
		return
	}
	// 打印日志，记录当前处理的模型及对应端口
	log.Printf("-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-")
	log.Printf("GGUF服务 -> 模型[ %s : %d ]", modelName, port)
	// 将请求反向代理到模型对应的端口
	proxyToPort(w, r, port)
}

// serveBusyResponse 返回“系统繁忙”的流式 SSE 响应
// 该函数用于在系统当前处理能力有限时，返回一个 SSE 格式的响应，提示用户系统当前正忙。
// 参数:
//   - w: http.ResponseWriter 类型，用于向客户端发送响应
func serveBusyResponse(w http.ResponseWriter) {
	// 设置响应头，指定内容类型为 SSE 格式，使用 UTF-8 编码
	w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	// 禁用缓存，确保每次请求都能获取最新响应
	w.Header().Set("Cache-Control", "no-cache")
	// 保持连接，确保 SSE 流式响应正常工作
	w.Header().Set("Connection", "keep-alive")
	// 写入 HTTP 状态码 200，表示请求成功
	w.WriteHeader(http.StatusOK)
	// 构造系统繁忙的响应数据
	response := map[string]any{
		"choices": []map[string]any{
			{
				// 完成原因设为 nil，表示响应未结束（这里只是格式要求）
				"finish_reason": nil,
				// 选择索引设为 0
				"index": 0,
				// 增量内容，提示用户系统繁忙
				"delta": map[string]any{
					"content": "请稍等哦, 月华现在正忙呢~~",
				},
			},
		},
		// 响应创建时间戳
		"created": time.Now().Unix(),
		// 响应 ID，添加当前纳秒时间戳确保唯一性
		"id": "chatcmpl-busy-" + fmt.Sprintf("%d", time.Now().UnixNano()),
		// 模型名称，标记为系统繁忙模型
		"model": "system-busy",
		// 系统指纹，添加当前纳秒时间戳确保唯一性
		"system_fingerprint": "busy-" + fmt.Sprintf("%d", time.Now().UnixNano()),
		// 对象类型，标记为聊天完成块
		"object": "chat.completion.chunk",
	}
	// 将响应数据编码为 JSON 格式
	jsonData, err := json.Marshal(response)
	// 若编码失败，返回 500 错误并提示生成响应失败
	if err != nil {
		http.Error(w, "GGUF服务[ERROR] -> 生成响应失败", http.StatusInternalServerError)
		return
	}
	// 按照 SSE 格式写入响应数据
	fmt.Fprintf(w, "data: %s\n\n", string(jsonData))
	// 写入 [DONE] 标记，表示 SSE 流结束
	fmt.Fprintf(w, "data: [DONE]\n\n")
	// 检查响应写入器是否支持 Flush 方法
	if flusher, ok := w.(http.Flusher); ok {
		// 刷新缓冲区，确保数据立即发送给客户端
		flusher.Flush()
	}
}

// extractModelName 从请求体或 URL 路径中提取模型名称
// 该函数根据 HTTP 请求的方法和路径，尝试提取模型名称。
// 如果请求方法为 POST，则从请求体中提取 "model" 字段的值；
// 如果请求方法为 GET，则从 URL 路径中提取模型名称。
// 参数:
//   - r: *http.Request 类型，代表 HTTP 请求
//
// 返回值:
//   - string: 提取到的模型名称，若提取失败则返回空字符串
func extractModelName(r *http.Request) string {
	// 用于存储从请求中提取的模型名称
	var modelName string
	// 用于存储请求体的字节数据，方便后续读取和恢复请求体
	var bodyBytes []byte
	// 检查请求方法是否为 POST，如果是则尝试从请求体中提取模型名称
	if r.Method == "POST" {
		var err error
		// 读取请求体的所有内容
		bodyBytes, err = io.ReadAll(r.Body)
		// 若读取失败，返回空字符串
		if err != nil {
			return ""
		}
		// 关闭请求体
		r.Body.Close()
		// 用于存储反序列化后的请求体数据
		var requestBody map[string]any
		// 将请求体内容反序列化为 map
		if err := json.Unmarshal(bodyBytes, &requestBody); err == nil {
			// 尝试从请求体中获取 "model" 字段的值
			if model, ok := requestBody["model"].(string); ok {
				modelName = model
			}
		}
		// 恢复请求体，确保后续处理可以再次读取请求体内容
		r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
		// 更新请求体长度
		r.ContentLength = int64(len(bodyBytes))
	}
	// 如果从请求体中未提取到模型名称，则尝试从 URL 路径中提取
	if modelName == "" {
		// 将 URL 路径按 "/" 分割成多个部分
		pathParts := strings.Split(r.URL.Path, "/")
		// 若分割后的路径部分数量大于 2，则取第 3 部分作为模型名称
		if len(pathParts) > 2 {
			modelName = pathParts[2]
		}
	}
	return modelName
}

// getModelPort 根据模型名称获取对应端口（加读锁）
// 该函数会从配置的模型端口映射中查找指定模型名称对应的端口号。
// 在访问模型端口映射时会加读锁，以确保并发安全，操作完成后自动解锁。
// 参数:
//   - modelName: 要查找的模型名称
//
// 返回值:
//   - int: 模型对应的端口号
//   - bool: 表示该模型是否存在于映射中
func getModelPort(modelName string) (int, bool) {
	// 加读锁，防止并发修改模型端口映射时出现数据竞争
	config.ModelMapMutex.RLock()
	// 函数结束时解锁，确保锁一定会被释放
	defer config.ModelMapMutex.RUnlock()
	// 从模型端口映射中查找指定模型的端口号
	port, exists := config.ModelPortMap[modelName]
	// 返回端口号和是否存在的标志
	return port, exists
}

// proxyToPort 将请求反向代理到指定本地端口
// 该函数接收一个 HTTP 响应写入器、HTTP 请求和端口号作为参数，
// 构造目标 URL 并解析，若解析成功则创建单主机反向代理并将请求转发到指定端口。
// 参数:
//   - w: http.ResponseWriter 类型，用于向客户端发送响应
//   - r: *http.Request 类型，代表客户端发送的 HTTP 请求
//   - port: int 类型，指定要反向代理到的本地端口号
func proxyToPort(w http.ResponseWriter, r *http.Request, port int) {
	// 构造目标 URL，格式为 http://localhost:端口号
	targetURL := fmt.Sprintf("http://localhost:%d", port)
	// 解析目标 URL，将字符串形式的 URL 转换为 url.URL 结构体
	target, err := url.Parse(targetURL)
	// 若解析失败，返回 500 错误并提示解析目标 URL 失败
	if err != nil {
		http.Error(w, "GGUF服务[ERROR] -> 解析目标 URL 失败", http.StatusInternalServerError)
		return
	}
	// 创建一个单主机反向代理，将请求转发到解析后的目标 URL
	proxy := httputil.NewSingleHostReverseProxy(target)
	// 执行反向代理，将请求转发到目标地址并将响应返回给客户端
	proxy.ServeHTTP(w, r)
}
